from typing import List, Dict
from instruction_re.core.datatypes import AnswerType, ContextFormat
from .Answer import AnswerFormatter
from .SentenceFormatter import SenteceFormatter
from .TypeExtractFormatter import TypeAnswerFormatter
from ..format_utils import (
    add_instruction_prefix,
    add_options_prefix,
    format_str_list_to_a_string,
)


class MainTaskFormatter:

    def __init__(
        self,
        answer_type: AnswerType = AnswerType.SVO,
        context_format: ContextFormat = ContextFormat.TAG,
        replace_label_to_special_token: bool = True,
        label2token: Dict[str, str] = None,
        re_strict: bool = False,
        sep: str = "<sep>",
        answer_order: int = 0,
        template_order: int = 3,
    ) -> None:
        """AnswerFormat的父类

        Args:
            answer_type (AnswerType, optional): _description_. Defaults to AnswerType.SVO.
            context_format (ContextFormat, optional): _description_. Defaults to ContextFormat.TAG.
            replace_label_to_special_token (bool, optional): _description_. Defaults to True.
            label2token (Dict[str, str], optional): _description_. Defaults to None.
            re_strict (bool, optional): _description_. Defaults to False.
            sep (str, optional): _description_. Defaults to "<sep>".
            answer_order (int, optional): 答案内部的输出顺序. Defaults to 0，
                对于[(John, work for, Microsoft), (Alice, kill, Bob)],
                answer_order = 0 输出 John is a person, Microsoft is an organization, Alice is a person, Bob is a person, John work for Microsoft, Alice kill Bob
                answer_order =1 输出 John is a person, Microsoft is an organization,John work for Microsoft,Alice is a person, Bob is a person,Alice kill Bob

        """
        self.answer_type = answer_type
        self.context_format = context_format
        self.replace_label_to_special_token = replace_label_to_special_token
        self.label2token = label2token
        self.sep = sep
        self.re_strict = re_strict
        # self.answer_order = 1
        self.template_order = template_order
        # 0: 先输出全部实体，再输出关系
        # 1: 先输出实体，再输出关系，按照实体的顺序输出
        # 2: 使用新的template 顺序输出 Havana (location) <organization_based_in> Radio Reloj Network (organization)

    def format_answer(self, relations: List[Dict[str, List[str]]]) -> str:
        answer_formatter = AnswerFormatter(
            use_special_token_for_label=self.replace_label_to_special_token,
            label2token=self.label2token,
            sep=self.sep,
            template_order=self.template_order,
        )
        if self.template_order == 0:
            return answer_formatter.fromat_relations_for_re_task(
                relations
            )

        elif self.template_order == 1:
            return answer_formatter.format_relation_quintuples_for_inside_order(
                relations
            )

        elif self.template_order == 2:
            return answer_formatter.format_relations_for_simple_template(relations)

        elif self.template_order == 3:
            return answer_formatter.from_relations_to_subject_verb_objext_structure(
                relations
            )

    def format_context(
        self,
        data_item: Dict,
    ) -> str:
        if self.re_strict or self.template_order == 3:
            return SenteceFormatter.format_sentence_for_re(data_item)
        elif self.context_format == ContextFormat.CONDITIONS.value:
            return SenteceFormatter.format_sentece_with_conditions(data_item=data_item)
        elif self.context_format == ContextFormat.TAG.value:
            return SenteceFormatter.format_sentence_with_tag(data_item=data_item)

    def format_type_answer(self, relations):
        return TypeAnswerFormatter.from_relations(
            relations, label2token=self.label2token
        )

    @classmethod
    def format_instruction_for_re_task(
        self,
        instruction: str,
        options: List[str],
    ) -> str:
        """为关系抽取任务格式化上下文，包括instruction, options

        Args:
            instruction: 任务的instruction
            options: 候选关系类型列表

        Returns:
            str: 格式化后的上下文
        """
        instruction_part = add_instruction_prefix(instruction)
        options_part = add_options_prefix(format_str_list_to_a_string(options))

        return instruction_part + " " + options_part


def main():

    sample = {
        "sentence": "John Wilkes Booth , who assassinated President Lincoln.",
        "relations": [
            {
                "head": {
                    "name": "John Wilkes Booth",
                    "type": "People",
                    "pos": [0, 17],
                },
                "type": "Kill",
                "tail": {
                    "name": "President Lincoln",
                    "type": "People",
                    "pos": [37, 54],
                },
            }
        ],
        "instruction": "Please type the relation between given entities based on the sentence.",
    }

    sample["examples"] = [
        {
            "sentence": "The opera company performed at the Palace of Fine Arts , in San Francisco , on June 30 and July 1-2 , said Kevin O 'Brien , a spokesman for the theater.",
            "relations": [
                {
                    "head": {
                        "name": "Palace of Fine Arts",
                        "type": "Location",
                        "pos": [35, 54],
                    },
                    "type": "Located in",
                    "tail": {
                        "name": "San Francisco",
                        "type": "Location",
                        "pos": [60, 73],
                    },
                }
            ],
        },
    ]

    formatter = MainTaskFormatter(answer_type=AnswerType.TEMPLATE)
    print(formatter.format_context(sample))


if __name__ == "__main__":
    main()
